package iscas;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.Pipeline;

import java.util.HashMap;
import java.util.function.Function;

@SuppressWarnings("unused")
public class RedisStateBackend {
    private final HashMap<String, HashMap<Object, Object>> map = new HashMap<>();
    private final HashMap<String, HashMap<Object, Object>> cache = new HashMap<>();
    private final Jedis jedis = RedisConfig.newJedis();
    private boolean fetch = false;

    public RedisStateBackend() {
    }

    public RedisStateBackend(boolean fetch) {
        this.fetch = fetch;
    }

    public Long getLong(String name, Object key) {
        return get(name, key, Long::parseLong);
    }

    public void setLong(String name, Object key, Long value) {
        set(name, key, value);
    }

    public Integer getInt(String name, Object key) {
        return get(name, key, Integer::parseInt);
    }

    public void setInt(String name, Object key, Integer value) {
        set(name, key, value);
    }

    private <T> T get(String name, Object key, Function<String, T> parser) {
        map.putIfAbsent(name, new HashMap<>());
        HashMap<Object, Object> m = map.get(name);
        m.computeIfAbsent(key, k -> {
            if (cache.containsKey(name)) {
                @SuppressWarnings("unchecked")
                T o = (T) cache.get(name).get(k);
                if (o != null) {
                    return o;
                }
            }
            if (fetch) {
                String r = jedis.hget(name, k.toString());
                return r == null ? null : parser.apply(r);
            }
            return null;
        });
        @SuppressWarnings("unchecked")
        T t = (T) m.get(key);
        return t;
    }

    private void set(String name, Object key, Object value) {
        map.putIfAbsent(name, new HashMap<>());
        map.get(name).put(key, value);
    }

    private void setCache(String name, Object key, Object value) {
        cache.putIfAbsent(name, new HashMap<>());
        cache.get(name).put(key, value);
    }

    public void close() {
        jedis.close();
    }

    public void sync() {
        Jedis jedis = RedisConfig.newJedis();
        Pipeline p = jedis.pipelined();
        map.forEach((String name, HashMap<Object, Object> m) ->
                m.forEach((Object k, Object v) -> p.hset(name, k.toString(), v.toString())));
        p.sync();
        map.clear();
        jedis.incr("syncCount");
        jedis.close();
    }

    public <T> void prefetch(String name, Function<String, Object> keyParser, Function<String, T> valueParser) {
        Jedis jedis = RedisConfig.newJedis();
        jedis.hgetAll(name).forEach((String k, String v) -> setCache(name, keyParser.apply(k), valueParser.apply(v)));
        jedis.incr("prefetchCount");
        jedis.close();
    }
}
